# -*- coding: utf-8 -*-
"""
   File Name:  load_dataset.py
   Author :    liccoo
   Time:       2022/8/24 12:43
"""
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader


class ForwardDataset(Dataset):  # 正设计 dataset
    def __init__(self, dataset_path):
        self.data = np.loadtxt(dataset_path, delimiter=',', skiprows=0, dtype=np.float32, encoding='utf-8-sig')
        self.x = self.data[:, 0:1]  # x
        self.y = self.data[:, 1:2]  # y = x^2

    def __getitem__(self, idx):
        return torch.tensor(self.x[idx], dtype=torch.float32), \
               torch.tensor(self.y[idx], dtype=torch.float32)

    def __len__(self):
        return len(self.data)


class InverseDataset(Dataset):  # 逆设计 dataset
    def __init__(self, dataset_path):
        self.data = np.loadtxt(dataset_path, delimiter=',', skiprows=0, dtype=np.float32, encoding='utf-8-sig')
        self.x = self.data[:, 0:1]  # x
        self.y = self.data[:, 1:2]  # y = x^2

    def __getitem__(self, idx):
        return torch.tensor(self.y[idx], dtype=torch.float32), \
               torch.tensor(self.x[idx], dtype=torch.float32)

    def __len__(self):
        return len(self.data)


def forward_dataloader(dataset_path, args):  # 正设计 dataloader
    dataset = ForwardDataset(dataset_path)
    return DataLoader(dataset=dataset, batch_size=args.batch_size, shuffle=args.shuffle, num_workers=args.num_workers)


def inverse_dataloader(dataset_path, args):  # 逆设计 dataloader
    dataset = InverseDataset(dataset_path)
    return DataLoader(dataset=dataset, batch_size=args.batch_size, shuffle=args.shuffle, num_workers=args.num_workers)
